{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ADST8G0Na4qU",
        "outputId": "77ddc936-dcb2-4459-b7ca-89d700a1ad14"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Overwriting test_conv2d_fp16.cu\n"
          ]
        }
      ],
      "source": [
        "%%writefile test_conv2d_fp16.cu\n",
        "#include <cuda_runtime.h>\n",
        "#include <cuda_fp16.h>\n",
        "#include <iostream>\n",
        "#include <random>\n",
        "\n",
        "// float32 kernel\n",
        "__global__ void conv2d_float(const float* input, const float* kernel, float* output,\n",
        "                            int input_rows, int input_cols,\n",
        "                            int kernel_rows, int kernel_cols,\n",
        "                            int output_rows, int output_cols) {\n",
        "    unsigned int out_col = blockIdx.x * blockDim.x + threadIdx.x;\n",
        "    unsigned int out_row = blockIdx.y * blockDim.y + threadIdx.y;\n",
        "\n",
        "    float sum = 0.0f;\n",
        "    if (out_row < output_rows && out_col < output_cols) {\n",
        "        for (int m = 0; m < kernel_rows; ++m) {\n",
        "            for (int n = 0; n < kernel_cols; ++n) {\n",
        "                int input_row = out_row + m;\n",
        "                int input_col = out_col + n;\n",
        "                float input_value = input[input_row * input_cols + input_col];\n",
        "                float kernel_value = kernel[m * kernel_cols + n];\n",
        "                sum += input_value * kernel_value;\n",
        "            }\n",
        "        }\n",
        "        output[out_row * output_cols + out_col] = sum;\n",
        "    }\n",
        "}\n",
        "\n",
        "// FP16 kernel =\n",
        "__global__ void conv2d_fp16(const __half* input, const __half* kernel,\n",
        "                          __half* output,\n",
        "                          int input_rows, int input_cols,\n",
        "                          int kernel_rows, int kernel_cols,\n",
        "                          int output_rows, int output_cols) {\n",
        "    unsigned int out_col = blockIdx.x * blockDim.x + threadIdx.x;\n",
        "    unsigned int out_row = blockIdx.y * blockDim.y + threadIdx.y;\n",
        "\n",
        "    __half sum = __float2half(0.0f);\n",
        "    if (out_row < output_rows && out_col < output_cols) {\n",
        "        for (int m = 0; m < kernel_rows; ++m) {\n",
        "            for (int n = 0; n < kernel_cols; ++n) {\n",
        "                int input_row = out_row + m;\n",
        "                int input_col = out_col + n;\n",
        "                __half input_value = input[input_row * input_cols + input_col];\n",
        "                __half kernel_value = kernel[m * kernel_cols + n];\n",
        "                sum = __hadd(sum, __hmul(input_value, kernel_value)); // Native FP16 ops\n",
        "            }\n",
        "        }\n",
        "        output[out_row * output_cols + out_col] = sum;\n",
        "    }\n",
        "}\n",
        "\n",
        "extern \"C\" float run_float_kernel(float* input, float* kernel, float* output,\n",
        "                                 int input_rows, int input_cols,\n",
        "                                 int kernel_rows, int kernel_cols) {\n",
        "    float *d_input, *d_kernel, *d_output;\n",
        "    int output_rows = input_rows - (kernel_rows - 1);\n",
        "    int output_cols = input_cols - (kernel_cols - 1);\n",
        "\n",
        "    cudaMalloc((void**)&d_input, input_rows * input_cols * sizeof(float));\n",
        "    cudaMalloc((void**)&d_kernel, kernel_rows * kernel_cols * sizeof(float));\n",
        "    cudaMalloc((void**)&d_output, output_rows * output_cols * sizeof(float));\n",
        "\n",
        "    cudaMemcpy(d_input, input, input_rows * input_cols * sizeof(float), cudaMemcpyHostToDevice);\n",
        "    cudaMemcpy(d_kernel, kernel, kernel_rows * kernel_cols * sizeof(float), cudaMemcpyHostToDevice);\n",
        "\n",
        "    dim3 threadsPerBlock(16, 16);\n",
        "    dim3 numBlocks((output_cols + threadsPerBlock.x - 1) / threadsPerBlock.x,\n",
        "                  (output_rows + threadsPerBlock.y - 1) / threadsPerBlock.y);\n",
        "\n",
        "    cudaEvent_t start, stop;\n",
        "    cudaEventCreate(&start);\n",
        "    cudaEventCreate(&stop);\n",
        "\n",
        "    cudaEventRecord(start);\n",
        "    conv2d_float<<<numBlocks, threadsPerBlock>>>(d_input, d_kernel, d_output,\n",
        "                                                input_rows, input_cols,\n",
        "                                                kernel_rows, kernel_cols,\n",
        "                                                output_rows, output_cols);\n",
        "    cudaEventRecord(stop);\n",
        "    cudaEventSynchronize(stop);\n",
        "\n",
        "    float milliseconds = 0;\n",
        "    cudaEventElapsedTime(&milliseconds, start, stop);\n",
        "\n",
        "    cudaMemcpy(output, d_output, output_rows * output_cols * sizeof(float), cudaMemcpyDeviceToHost);\n",
        "\n",
        "    cudaFree(d_input); cudaFree(d_kernel); cudaFree(d_output);\n",
        "    cudaEventDestroy(start); cudaEventDestroy(stop);\n",
        "\n",
        "    return milliseconds / 1000.0f; // Convert to seconds\n",
        "}\n",
        "\n",
        "extern \"C\" float run_fp16_kernel(float* input, float* kernel, float* output,\n",
        "                                int input_rows, int input_cols,\n",
        "                                int kernel_rows, int kernel_cols) {\n",
        "    __half *d_input, *d_kernel, *d_output;\n",
        "    int output_rows = input_rows - (kernel_rows - 1);\n",
        "    int output_cols = input_cols - (kernel_cols - 1);\n",
        "\n",
        "    // Convert to FP16 once on host\n",
        "    __half *h_input_fp16 = new __half[input_rows * input_cols];\n",
        "    __half *h_kernel_fp16 = new __half[kernel_rows * kernel_cols];\n",
        "    for (int i = 0; i < input_rows * input_cols; i++)\n",
        "        h_input_fp16[i] = __float2half(input[i]);\n",
        "    for (int i = 0; i < kernel_rows * kernel_cols; i++)\n",
        "        h_kernel_fp16[i] = __float2half(kernel[i]);\n",
        "\n",
        "    cudaMalloc((void**)&d_input, input_rows * input_cols * sizeof(__half));\n",
        "    cudaMalloc((void**)&d_kernel, kernel_rows * kernel_cols * sizeof(__half));\n",
        "    cudaMalloc((void**)&d_output, output_rows * output_cols * sizeof(__half));\n",
        "\n",
        "    cudaMemcpy(d_input, h_input_fp16, input_rows * input_cols * sizeof(__half), cudaMemcpyHostToDevice);\n",
        "    cudaMemcpy(d_kernel, h_kernel_fp16, kernel_rows * kernel_cols * sizeof(__half), cudaMemcpyHostToDevice);\n",
        "\n",
        "    dim3 threadsPerBlock(16, 16);\n",
        "    dim3 numBlocks((output_cols + threadsPerBlock.x - 1) / threadsPerBlock.x,\n",
        "                  (output_rows + threadsPerBlock.y - 1) / threadsPerBlock.y);\n",
        "\n",
        "    cudaEvent_t start, stop;\n",
        "    cudaEventCreate(&start);\n",
        "    cudaEventCreate(&stop);\n",
        "\n",
        "    cudaEventRecord(start);\n",
        "    conv2d_fp16<<<numBlocks, threadsPerBlock>>>(d_input, d_kernel, d_output,\n",
        "                                               input_rows, input_cols,\n",
        "                                               kernel_rows, kernel_cols,\n",
        "                                               output_rows, output_cols);\n",
        "    cudaEventRecord(stop);\n",
        "    cudaEventSynchronize(stop);\n",
        "\n",
        "    float milliseconds = 0;\n",
        "    cudaEventElapsedTime(&milliseconds, start, stop);\n",
        "\n",
        "    __half *h_output_fp16 = new __half[output_rows * output_cols];\n",
        "    cudaMemcpy(h_output_fp16, d_output, output_rows * output_cols * sizeof(__half), cudaMemcpyDeviceToHost);\n",
        "\n",
        "    for (int i = 0; i < output_rows * output_cols; i++)\n",
        "        output[i] = __half2float(h_output_fp16[i]);\n",
        "\n",
        "    cudaFree(d_input); cudaFree(d_kernel); cudaFree(d_output);\n",
        "    delete[] h_input_fp16; delete[] h_kernel_fp16; delete[] h_output_fp16;\n",
        "    cudaEventDestroy(start); cudaEventDestroy(stop);\n",
        "\n",
        "    return milliseconds / 1000.0f; // Convert to seconds\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# compile the CUDA code into a shared library\n",
        "!nvcc -shared -o test_conv2d_fp16.so test_conv2d_fp16.cu -arch=sm_75 -Xcompiler -fPIC"
      ],
      "metadata": {
        "id": "u6Ab5nS6bVha"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "from ctypes import *\n",
        "import time\n",
        "\n",
        "lib = cdll.LoadLibrary('./test_conv2d_fp16.so')\n",
        "\n",
        "# update function signatures to return float\n",
        "lib.run_float_kernel.argtypes = [POINTER(c_float), POINTER(c_float), POINTER(c_float),\n",
        "                                c_int, c_int, c_int, c_int]\n",
        "lib.run_float_kernel.restype = c_float\n",
        "lib.run_fp16_kernel.argtypes = [POINTER(c_float), POINTER(c_float), POINTER(c_float),\n",
        "                               c_int, c_int, c_int, c_int]\n",
        "lib.run_fp16_kernel.restype = c_float\n",
        "\n",
        "# test data\n",
        "input_rows, input_cols = 1024 * 3, 1024 * 3\n",
        "kernel_rows, kernel_cols = 128, 128\n",
        "output_rows = input_rows - (kernel_rows - 1)\n",
        "output_cols = input_cols - (kernel_cols - 1)\n",
        "\n",
        "input_data = np.random.randn(input_rows, input_cols).astype(np.float32)\n",
        "kernel_data = np.random.randn(kernel_rows, kernel_cols).astype(np.float32)\n",
        "output_float = np.zeros((output_rows, output_cols), dtype=np.float32)\n",
        "output_fp16 = np.zeros((output_rows, output_cols), dtype=np.float32)\n",
        "\n",
        "input_ptr = input_data.ctypes.data_as(POINTER(c_float))\n",
        "kernel_ptr = kernel_data.ctypes.data_as(POINTER(c_float))\n",
        "output_float_ptr = output_float.ctypes.data_as(POINTER(c_float))\n",
        "output_fp16_ptr = output_fp16.ctypes.data_as(POINTER(c_float))\n",
        "\n",
        "num_runs = 100\n",
        "warmup_runs = 10\n",
        "\n",
        "float_times = []\n",
        "fp16_times = []\n",
        "\n",
        "# warm-up\n",
        "for _ in range(warmup_runs):\n",
        "    lib.run_float_kernel(input_ptr, kernel_ptr, output_float_ptr,\n",
        "                        input_rows, input_cols, kernel_rows, kernel_cols)\n",
        "    lib.run_fp16_kernel(input_ptr, kernel_ptr, output_fp16_ptr,\n",
        "                       input_rows, input_cols, kernel_rows, kernel_cols)\n",
        "\n",
        "# benchmark\n",
        "for _ in range(num_runs):\n",
        "    float_time = lib.run_float_kernel(input_ptr, kernel_ptr, output_float_ptr,\n",
        "                                     input_rows, input_cols, kernel_rows, kernel_cols)\n",
        "    float_times.append(float_time)\n",
        "\n",
        "    fp16_time = lib.run_fp16_kernel(input_ptr, kernel_ptr, output_fp16_ptr,\n",
        "                                   input_rows, input_cols, kernel_rows, kernel_cols)\n",
        "    fp16_times.append(fp16_time)\n",
        "\n",
        "# stats\n",
        "float_mean = np.mean(float_times)\n",
        "float_std = np.std(float_times)\n",
        "fp16_mean = np.mean(fp16_times)\n",
        "fp16_std = np.std(fp16_times)\n",
        "\n",
        "diff = np.abs(output_float - output_fp16)\n",
        "max_error = np.max(diff)\n",
        "mean_error = np.mean(diff)\n",
        "mse = np.mean(diff**2)\n",
        "\n",
        "print(f\"Float32 kernel time: {float_mean:.6f} ± {float_std:.6f} seconds ({num_runs} runs)\")\n",
        "print(f\"FP16 kernel time: {fp16_mean:.6f} ± {fp16_std:.6f} seconds ({num_runs} runs)\")\n",
        "print(f\"Speedup (Float32/FP16): {float_mean/fp16_mean:.2f}x\")\n",
        "print(f\"Maximum absolute error: {max_error:.6f}\")\n",
        "print(f\"Mean absolute error: {mean_error:.6f}\")\n",
        "print(f\"Mean squared error: {mse:.6f}\")\n",
        "\n",
        "print(\"\\nFirst 5x5 of Float32 output:\")\n",
        "print(output_float[:5, :5])\n",
        "print(\"\\nFirst 5x5 of FP16 output:\")\n",
        "print(output_fp16[:5, :5])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KqstzU0ibNHm",
        "outputId": "a2d7ba7f-e1a8-4432-b98e-11aeec429e07"
      },
      "execution_count": 26,
      "outputs": [
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Float32 kernel time: 0.000000 ± 0.000000 seconds (100 runs)\n",
            "FP16 kernel time: 0.000000 ± 0.000000 seconds (100 runs)\n",
            "Speedup (Float32/FP16): 1.46x\n",
            "Maximum absolute error: 0.289307\n",
            "Mean absolute error: 0.036147\n",
            "Mean squared error: 0.002109\n",
            "\n",
            "First 5x5 of Float32 output:\n",
            "[[ -55.84884     25.658815  -413.50046    259.5033    -109.114334 ]\n",
            " [ -58.657173    59.219414   218.30032     57.322186    -8.62942  ]\n",
            " [-166.65118    214.80194    314.93027   -192.7418    -111.4903   ]\n",
            " [ 186.23802   -160.4699    -162.89262    -34.438816    17.689949 ]\n",
            " [ 302.5674      -6.6662626  217.40765     16.765682   249.54005  ]]\n",
            "\n",
            "First 5x5 of FP16 output:\n",
            "[[ -55.84375     25.671875  -413.5        259.5       -109.0625   ]\n",
            " [ -58.6875      59.21875    218.25        57.25        -8.6171875]\n",
            " [-166.625      214.75       315.        -192.75      -111.5      ]\n",
            " [ 186.25      -160.5       -162.875      -34.4375      17.6875   ]\n",
            " [ 302.5         -6.703125   217.375       16.8125     249.5      ]]\n"
          ]
        }
      ]
    }
  ]
}